#pragma once

#include "./common.h"
#include "./util.hpp"
#include <unordered_map>

#include "./thread_pool.hpp"
#include "weight.hpp"

struct DocInfo{
    std::string     title;
    std::string     content;
    std::string     url;
    std::uint64_t   id;
    //
    std::uint32_t   length = 0;//分词总数
};

struct InvertedElement{
    std::uint64_t id;       //
    std::string word;       //jieba 分词分出来的  注意这里并不冗余哦
    //因为 InvertedElement 是最终要以数组的形式 我们获得一个 DocInfo 对应的 vector<InvertedList>
    //最后插入到 inverted_index 要遍历 vector<InvertedList> 插入到对应的 inverted_index 的 key。这时 key 就要 word 了
    double weight;
    // update
    int title_cnt = 0;
    int content_cnt = 0;
};

// InvertedElement for deduplication
struct InvertedElement_dedup{
    std::uint64_t id;
    double weight;
    std::vector<std::string> word;
    InvertedElement_dedup() : weight(0.0),id(-1){}

    void toStr(){
        std::cout << id << " " << word << " " << weight << "\n";
    }
};
// 重要组件，正倒排索引
class Index{
    private:
        // 规定下标就是 Doc 的 ID
        using Inverted_list_t = std::vector<InvertedElement>;
        std::vector<DocInfo>                                forward_index;
        std::unordered_map<std::string, Inverted_list_t>    inverted_index;

        double avg_doc_len = 0.0;

    public:
        static Index *getInstance(){
            //静态类成员不安全，而且instance建议用成懒汉模式
            static Index instance;
            return &instance;
        }
    public:
        DocInfo *getDocInfo(uint64_t id){
            if(id >= forward_index.size()){
                std::cerr << "get DocInfo failed. ID: " << id << "\n";
                return nullptr;
            }
            return &forward_index[id];
        }
        Inverted_list_t *getInvertedList(std::string key){
            auto iter = inverted_index.find(key);
            if(iter == inverted_index.end()) {
                DBG_LOG("get InvertedList ERROR: key %s NOT FOUND!\n", key.c_str());
                return nullptr;
            }
            return &(iter->second);
        }
    public:
        // 正倒排持久化保存
        //                            std::vector<DocInfo>    forward_index;
        //std::unordered_map<std::string, Inverted_list_t>    inverted_index;

        //反序列化  Deserialize
        bool LoadForwardIndex(const std::string &path) {
            std::ifstream in(path, std::ios::binary);
            if(!in.is_open()){
                ERROR_LOG("LoadForwardIndex Error: %s Open Failed\n", path.c_str());
                return false;
            }
            uint32_t magic, version;
            in.read((char *)&magic, sizeof(magic));
            if(magic != forward_magic) {
                ERROR_LOG("LoadForwardIndex Error: Magic num mismatch\n");
                return false;
            }
            in.read((char*)&version, sizeof(version));
            if (version != 1) {
                ERROR_LOG("LoadForwardIndex: version %u not supported\n", version);
                return false;
            }
            uint64_t doc_nr;
            in.read((char*)&doc_nr, sizeof(doc_nr));
            forward_index.clear();
            forward_index.reserve(doc_nr);
            
            auto read_string = [&](std::string &s){
                uint64_t sz;
                in.read((char*)&sz, sizeof(sz));
                s.resize(sz);
                if(sz > 0) in.read(s.data(), sz);
            };
            for(uint64_t i = 0; i < doc_nr; ++i){
                DocInfo doc;
                in.read((char*)&doc.id, sizeof(doc.id));
                in.read((char*)&doc.length, sizeof(doc.length));
                
                // 读 string
                read_string(doc.title);
                read_string(doc.content);
                read_string(doc.url);
                forward_index.emplace_back(std::move(doc));
            }
            NORM_LOG("LoadForwardIndex: load %llu docs from %s\n", (unsigned long long)doc_nr, path.c_str());
            return true;
        }
        bool LoadInvertedIndex(const std::string &path) {
            std::ifstream in(path, std::ios::binary);
            if(!in.is_open()){
                ERROR_LOG("LoadInvertedIndex Error: %s Open Failed\n", path.c_str());
                return false;
            }

            uint32_t magic, version;
            in.read((char*)&magic, sizeof(magic));
            if(magic != inverted_magic){
                ERROR_LOG("LoadInvertedIndex Error: Magic num mismatch\n");
                return false;
            }
            in.read((char*)&version, sizeof(version));
            if(version != 1){
                ERROR_LOG("LoadForwardIndex: version %u not supported\n", version);
                return false;
            }

            uint64_t term_nr;
            in.read((char*)&term_nr, sizeof(term_nr));

            auto read_string = [&](std::string &s){
                uint64_t sz;
                in.read((char*)&sz, sizeof(sz));
                s.resize(sz);
                if(sz > 0) in.read(s.data(), sz);
            };

            for(uint64_t i = 0; i < term_nr; ++i){
                std::string word;
                read_string(word);
                uint64_t list_sz;
                in.read((char*)&list_sz, sizeof(list_sz));
                Inverted_list_t list;
                list.resize(list_sz);
                for(uint64_t j = 0; j < list_sz; ++j){
                    auto &elem = list[j];
                    in.read((char*)&elem.id,            sizeof(elem.id));
                    in.read((char*)&elem.title_cnt,     sizeof(elem.title_cnt));
                    in.read((char*)&elem.content_cnt,   sizeof(elem.content_cnt));
                    in.read((char*)&elem.weight,        sizeof(elem.weight));
                    elem.word = word;
                }
                inverted_index.emplace(std::move(word), std::move(list));
                // inverted_index[word] = std::move(list);
            }
            NORM_LOG("LoadInvertedIndex: load %llu terms from %s\n", (unsigned long long)term_nr, path.c_str());

            return true;
        }
        //序列化    Serialize
        bool SaveForwardIndex(const std::string &path) const {
            std::ofstream out(path, std::ios::binary);
            if(!out.is_open()) {
                ERROR_LOG("SaveForwardIndex: open %s failed\n", path.c_str());
                return false;
            }
            const uint32_t magic = forward_magic;//随便写的一个
            const uint32_t version = 1;

            out.write((const char*)&magic, sizeof(magic));
            out.write((const char*)&version, sizeof(version));

            uint64_t doc_count = forward_index.size();
            out.write((const char*)&doc_count, sizeof(doc_count));

            auto write_string = [&](const std::string &s){
                uint64_t sz = s.size();
                out.write((const char*)&sz, sizeof(sz));
                out.write(s.data(), sz);
            };

            for(const auto &doc : forward_index) {
                // Save DocInfo
                out.write((const char*)&doc.id, sizeof(doc.id));
                out.write((const char*)&doc.length, sizeof(doc.length));

                write_string(doc.title);
                write_string(doc.content);
                write_string(doc.url);
            }
            NORM_LOG("SaveForwardIndex: save %llu terms to %s\n", (unsigned long long)doc_count, path.c_str());
            return true;
        }
        bool SaveInvertedIndex(const std::string &path) const {
            std::ofstream out(path, std::ios::binary);
            if(!out.is_open()){
                ERROR_LOG("SaveInvertedIndex: open %s failed\n", path.c_str());
                return false;
            }
            const uint32_t magic = inverted_magic;
            const uint32_t version = 1;
            
            out.write((const char*)&magic, sizeof(magic));
            out.write((const char*)&version, sizeof(version));

            uint64_t term_cnt = inverted_index.size();
            out.write((const char*)&term_cnt, sizeof(term_cnt));

            auto write_string = [&](const std::string &s){
                uint64_t sz = s.size();
                out.write((const char*)&sz, sizeof(sz));
                out.write(s.data(), sz);
            };

            for(const auto&[word, inverted_list] : inverted_index){
                // key 先写入
                write_string(word);
                uint64_t list_sz = inverted_list.size();
                out.write((const char*)&list_sz, sizeof(list_sz));

                for(const auto &elem : inverted_list){
                    out.write((const char*)&elem.id,            sizeof(elem.id)); 
                    out.write((const char*)&elem.title_cnt,     sizeof(elem.title_cnt)); 
                    out.write((const char*)&elem.content_cnt,   sizeof(elem.content_cnt)); 
                    out.write((const char*)&elem.weight,        sizeof(elem.weight)); 
                    // elem.word 不用写，反序列化时本身就有，还剩空间
                }
            }
            DBG_LOG("SaveInvertedIndex: save %llu terms to %s\n", (unsigned long long)term_cnt, path.c_str());
            return true;
        }
    private:
        // 这个函数在 构建正倒排索引后调用，统一对文档进行权重设置
        void ApplyWeightStrategy(const IWeightStrategy &strategy){
            const std::uint64_t N = forward_index.size();// 文档数
            if(N == 0) return;

            std::uint64_t sum_len = 0;
            for(const auto & doc : forward_index) sum_len += doc.length;
            avg_doc_len = static_cast<double>(sum_len) / static_cast<double>(N);

            // 遍历每个分词 term 的倒排表，计算weight
            for(auto &kv : inverted_index){
                const std::string &term = kv.first;
                Inverted_list_t &inverted_list = kv.second;

                CorpusStat corpus_stat;
                corpus_stat.N = N;
                corpus_stat.df = inverted_list.size();
                corpus_stat.avg_doc_len = avg_doc_len;

                for(auto &elem : inverted_list) {
                    TermLocalStat term_stat;
                    term_stat.content_cnt = elem.content_cnt;
                    term_stat.title_cnt   = elem.title_cnt;

                    const DocInfo &doc = forward_index[elem.id];
                    DocStat doc_stat;
                    doc_stat.length = doc.length;
                    elem.weight = strategy.compute(term_stat, doc_stat, corpus_stat);
                }
            }
        }
    private:
        /**
         * @note 根据构造时我们的协议: three part was apart by Seprator:'\3', and between 2 doc there is a new line char '\n'
         */
        DocInfo *BuildForwardIndex(const std::string &line){
            std::vector<std::string> result;
            util_ns::StringUtil::CutString(line, separator, &result);
            if(result.size() != 3){
                return nullptr;
            }
            DocInfo tmp;
            tmp.title = result[0];
            tmp.content = result[1];
            tmp.url = result[2];
            tmp.id = forward_index.size();
            forward_index.push_back(std::move(tmp));
            return &forward_index.back();
        }

        std::vector<InvertedElement> MakeInvertedElements(DocInfo &doc){
            struct word_cnt{
                int title_cnt = 0;
                int content_cnt = 0;
            };
            std::unordered_map<std::string, word_cnt> wordAndCnts;
            
            // title
            std::vector<std::string> title_words;
            util_ns::JiebaUtil::WordSplit(doc.title, &title_words);
            for(std::string word:title_words){
                boost::to_lower(word);
                ++wordAndCnts[word].title_cnt;
            }
            // content
            std::vector<std::string> content_words;
            util_ns::JiebaUtil::WordSplit(doc.content, &content_words);
            for(std::string word:content_words){
                boost::to_lower(word);
                ++wordAndCnts[word].content_cnt;
            }
            for(auto &[word, cnt] : wordAndCnts){
                doc.length += cnt.content_cnt + cnt.title_cnt;
            }
            // 
            std::vector<InvertedElement> elems;
            elems.reserve(wordAndCnts.size());
            for(std::pair<const std::string, word_cnt> &P_wordcnt: wordAndCnts){
                InvertedElement elem;
                elem.id = doc.id;
                elem.word = P_wordcnt.first;
                // update
                elem.content_cnt = P_wordcnt.second.content_cnt;
                elem.title_cnt = P_wordcnt.second.title_cnt;
                // weight 设置
                // elem.weight = static_cast<double>(P_wordcnt.second.content_cnt + P_wordcnt.second.title_cnt * 10);
                // 现在 weight 留白，等之后排序前我们统一统计一下所有的参数统一赋权值
                elems.push_back(std::move(elem));
            }
            return elems;
        }
        // 单线程建立倒排
        void BuildInvertedIndex(DocInfo &doc){
            auto elems = MakeInvertedElements(doc);
            for(auto &elem : elems){
                inverted_index[elem.word].push_back(std::move(elem));
            }
        }
    public:
        // 串行正排，多线程倒排
        void BuildIndex(const std::string &input_path){
            std::ifstream in(input_path, std::ios::in | std::ios::binary);
            if(!in.is_open()){
                DBG_LOG("BuildIndex: open failed\n");
                return;
            }
            std::string line;
            int forward_cnt = 0;
            //正排进度输出控制
            const static int Forward_Print_Gap = 1000;
            while(getline(in, line)){
                if((forward_cnt % Forward_Print_Gap) == 0) NORM_LOG("Built ForwardIndex item: %d\n", forward_cnt);
                if(line.empty()){
                    std::cerr << "BuildIndex: getline Error continue\n";
                    continue;
                }
                // ---正排---
                DocInfo *doc = BuildForwardIndex(line);
                if(!doc) continue;
                ++forward_cnt;
            }
            NORM_LOG("Built ForwardIndex item Done. total item: %d\n", forward_cnt);
            // ---多线程 建倒排---
            //倒排进度输出控制
            // const static int Inverted_Print_Gap = 200;
            int submitted = 0;                  //主线程提交的
            std::atomic<int> finished{0};    //任务线程完成的,直接原子类就行了
            const std::size_t n_docs = forward_index.size();
            if(n_docs <= 0) return;
            std::size_t n_threads = std::thread::hardware_concurrency();
            if(n_threads == 0) n_threads = 4;
            // ThreadPool pool(n_threads);
            ThreadPool &pool = ThreadPool::getInstance(n_threads);
            std::mutex mtx;//lock inverted_index(写时)
            std::vector<std::future<void>> futures;
            for(std::size_t i = 0; i < n_docs; ++i){
                DocInfo *doc = &forward_index[i];
                // if((submitted % Inverted_Print_Gap) == 0) NORM_LOG("Submitted BuildInvertedIndex task: %d\n", submitted);

                futures.emplace_back(
                    //用一个mtx 在写入时用,doc这里要值拷贝进去，因为是异步的，doc生命周期可能在调用lambda时就没了
                    pool.submit([this, doc, &mtx, &finished](){
                        //和BuildInvertedIndex（单线程版本）类似只是这里多线程要加锁
                        auto elems = MakeInvertedElements(*doc);
                        {
                            std::lock_guard<std::mutex> lk(mtx);
                            for(auto &elem : elems){
                                inverted_index[elem.word].push_back(std::move(elem));
                            }
                        }
                        int done = finished.fetch_add(1) + 1;
                        // if (done % 1000 == 0) NORM_LOG("InvertedIndex finished: %d\n", done);
                    })
                );
                ++submitted;
            }
            NORM_LOG("Submitted BuildInvertedIndex task: %d\n",submitted);
            // === 进度条 ===
            auto print_progress = [&](int done){
                const int bar_width = 50; // 进度条宽度
                double ratio = static_cast<double>(done) / static_cast<double>(n_docs);
                if(ratio > 1.0) ratio = 1.0;

                int pos = static_cast<int>(ratio * bar_width);

                std::cout << "\r[";
                for(int i = 0; i < bar_width; ++i){
                    if(i < pos)         std::cout << "=";
                    else if(i == pos)   std::cout << ">";
                    else                std::cout << " ";
                }
                std::cout << "] ";
                std::cout << static_cast<int>(ratio * 100.0) << "% (" 
                        << done << "/" << n_docs << ")";
                std::cout.flush();
            };

            // 主线程轮询 finished，刷新进度条
            while(true){
                int done = finished.load(std::memory_order_relaxed);
                print_progress(done);
                if(done >= static_cast<int>(n_docs)) break;
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
            }
            std::cout << std::endl; // 进度条完成后换行
            
            for(auto &f:futures){
                f.get();
            }
            int final_finished = finished.load();
            NORM_LOG("Built InvertedIndex item Done. total item: %d\n", final_finished);

            NORM_LOG("BuildIndex: Build InvertedIndex done with %zu docs, %zu threads\n", n_docs, n_threads);
            TFIDFStrategy strategy;
            // BM25Strategy strategy;
            ApplyWeightStrategy(strategy);
        }

};
